import os
import sys
import re
from typing import Dict, Any

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import OPTIMIZER_PROMPT
from agents.base_agent import BaseAgent

class OptimizerAgent(BaseAgent):
    def __init__(self):
        super().__init__("优化者")
        
        # 加载提示词
        with open(OPTIMIZER_PROMPT, 'r', encoding='utf-8') as f:
            optimizer_template = f.read()
        
        # 创建提示模板
        self.prompt = PromptTemplate(
            template=optimizer_template,
            input_variables=[
                "current_documentation", 
                "improvement_suggestions",
                "iteration",
                "task_type",
                "task_name",
                "optimization_focus"
            ]
        )
    
    def optimize_documentation(self, current_documentation, improvement_suggestions=None, iteration=1, task_info=None):
        """
        评审文档并提供反馈。
        
        参数:
            current_documentation (str): 当前需要优化的文档内容。
            improvement_suggestions (str 或 list): 改进建议。
            iteration (int): 当前迭代次数。
            task_info (dict): 任务相关信息，可包含：
                - type: 任务类型（class/function）
                - name: 任务名称
                - class_name: 类名（如适用）
                - pending_issues: 待解决的问题列表
                - iteration_history: 优化历史记录
        """
        # 整合改进建议
        if isinstance(improvement_suggestions, str):
            improvement_suggestions = improvement_suggestions.strip()
        elif isinstance(improvement_suggestions, list):
            improvement_suggestions = "\n".join([s.strip() for s in improvement_suggestions if s.strip()])
        else:
            improvement_suggestions = ""
        
        # 获取基本任务信息
        task_type = task_info.get('type', '') if task_info else ''
        task_name = task_info.get('name', '') if task_info else ''
        class_name = task_info.get('class_name', '') if task_info else ''
        
        # 获取优化历史和待解决问题
        pending_issues = task_info.get('pending_issues', []) if task_info else []
        optimization_history = task_info.get('iteration_history', []) if task_info else []
        
        # 任务描述美化
        if task_type == "class":
            task_description = f"{task_type} '{task_name}'"
        elif class_name:
            task_description = f"方法 '{task_name}' (类 '{class_name}' 的成员)"
        else:
            task_description = f"函数 '{task_name}'"
        
        # 构建历史上下文
        history_context = ""
        if optimization_history:
            history_items = []
            for history in optimization_history:
                improvements = ", ".join(history['key_improvements']) if history['key_improvements'] else "一般改进"
                history_items.append(f"- 迭代 {history['iteration']}: 评分 {history['score']}, 变化 {history['score_change']:+.1f}, 主要改进: {improvements}")
            
            history_context = "优化历史:\n" + "\n".join(history_items)
        
        # 待解决问题上下文
        issues_context = ""
        if pending_issues:
            issues_context = "需要解决的问题:\n- " + "\n- ".join(pending_issues)
        
        # 构建优化重点
        if iteration == 1:
            optimization_focus = f"初始文档生成，确保基本结构完整，覆盖{task_description}的主要功能和用途。"
        elif pending_issues:
            optimization_focus = f"重点解决上述问题，同时提升文档的整体专业度和可读性。"
        elif iteration == 2:
            optimization_focus = f"增强技术细节和准确性，添加更详细的参数说明和返回值解释。"
        else:
            optimization_focus = f"提炼语言表达，确保文档简洁、精准，符合专业技术文档标准。"

        # 准备输入数据
        inputs = {
            'current_documentation': current_documentation,
            'improvement_suggestions': improvement_suggestions,
            'iteration': str(iteration),
            'task_type': task_type,
            'task_name': task_name,
            'optimization_focus': optimization_focus,
            'issues_context': issues_context,
            'history_context': history_context
        }

        # 使用流式输出运行优化过程
        optimize_result = self.run_with_streaming(
            self.prompt,
            inputs,
            step=f"优化文档 ({task_description}, 迭代 {iteration})"
        )

        # 返回优化结果
        return {
            'documentation': optimize_result,
            'inputs': {
                'documentation_length': len(current_documentation),
                'pending_issues_count': len(pending_issues),
                'iteration': iteration
            }
        }
